# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score, classification_report
import xgboost as xgb
from json_repair import repair_json
import json_repair
import pickle
import joblib  # 用于模型保存

class Run_xgb:
    def __init__(self,):
        self.model = None
        self.save_path = None
        self.le = None
        # 7. 设置XGBoost参数
        self.params = {
            'objective': 'multi:softmax',  # 多分类目标
            'num_class': 4,  # 类别数量
            'eval_metric': 'mlogloss',  # 多分类对数损失
            'eta': 0.1,  # 学习率
            'max_depth': 6,  # 树深度
            'subsample': 0.8,  # 样本采样率
            'colsample_bytree': 0.8,  # 特征采样率
            'seed': 42,
            'nthread': 4
        }
        self.tfidf = None
    def set_save_path(self, path):
        self.save_path = path
    def get_save_path(self):
        return self.save_path
    def train_on_file(self, filename):
        # 1. 加载数据集（假设CSV包含text列和label列）
        # 示例数据格式：
        # text,label
        # "This is a sports news...", sports
        # "Technology breakthrough...", tech
        # ...共四种类别
        with open(filename, 'r') as f:
            topic_list = json_repair.load(f)
        df = pd.DataFrame(topic_list)
        df['cleaned_question'] = df['question'].apply(self.clean_text)
        # 3. 标签编码
        self.le = LabelEncoder()
        df['label_encoded'] = self.le.fit_transform(df['label'])
        # 4. 特征提取 - TF-IDF
        self.tfidf = TfidfVectorizer(
            max_features=5000,  # 根据内存调整
            ngram_range=(1, 2),  # 包含unigram和bigram
            stop_words='english'  # 移除英文停用词
        )

        X = self.tfidf.fit_transform(df['cleaned_question'])
        y = df['label_encoded']

        # 5. 划分数据集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42, stratify=y
        )

        # 6. 转换为DMatrix格式（优化内存）
        dtrain = xgb.DMatrix(X_train, label=y_train)
        dtest = xgb.DMatrix(X_test, label=y_test)

        # 8. 训练模型
        self.model = xgb.train(
            self.params,
            dtrain,
            num_boost_round=200,
            evals=[(dtrain, 'train'), (dtest, 'test')],
            early_stopping_rounds=20,
            verbose_eval=10
        )

        # 9. 预测与评估
        y_pred = self.model.predict(dtest)
        print(f"Accuracy: {accuracy_score(y_test, y_pred):.4f}")

    @staticmethod
    def clean_text(text:str)->str:
        text = text.lower()  # 转小写
        text = ''.join([c for c in text if c.isalnum() or c.isspace()])  # 移除非字母数字字符
        return text
    def predict(self, text):
        cleaned_text = self.clean_text(text)
        #tfidf = TfidfVectorizer(vocabulary=self.tfidf.vocabulary_)
        X = self.tfidf.transform([cleaned_text])
        y_pred = self.model.predict(xgb.DMatrix(X))
        predicted_label = int(y_pred[0]+1)
        return predicted_label

if __name__ == "__main__":
    run_xgb = Run_xgb()
    run_xgb.train_on_file("..\\..\\topic_list.json")

    with open("..\\..\\xgb_model_output\\xgb_model.data","wb") as file:
        pickle.dump(run_xgb, file)
